from data_processing import *
from utils.data_load import dataload
import matplotlib.pyplot as plt
# 导入评估指标
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score
from xgboost import XGBClassifier
import joblib
import os

# 数据加载
X_train, X_test, y_train, y_test = dataload()

# 特征构造
# --- 比例类特征 ---
X_train['YearsAtCompany_Age_ratio'] = X_train['YearsAtCompany'] / (X_train['Age'] + 1e-5)
X_test['YearsAtCompany_Age_ratio'] = X_test['YearsAtCompany'] / (X_test['Age'] + 1e-5)

X_train['TotalWorkingYears_Age_ratio'] = X_train['TotalWorkingYears'] / (X_train['Age'] + 1e-5)
X_test['TotalWorkingYears_Age_ratio'] = X_test['TotalWorkingYears'] / (X_test['Age'] + 1e-5)

X_train['CurrentRole_Company_ratio'] = X_train['YearsInCurrentRole'] / (X_train['YearsAtCompany'] + 1e-5)
X_test['CurrentRole_Company_ratio'] = X_test['YearsInCurrentRole'] / (X_test['YearsAtCompany'] + 1e-5)

# --- 收入相关 ---
X_train['Income_per_Year'] = X_train['MonthlyIncome'] / (X_train['TotalWorkingYears'] + 1)
X_test['Income_per_Year'] = X_test['MonthlyIncome'] / (X_test['TotalWorkingYears'] + 1)

X_train['Income_per_Level'] = X_train['MonthlyIncome'] / (X_train['JobLevel'] + 1e-5)
X_test['Income_per_Level'] = X_test['MonthlyIncome'] / (X_test['JobLevel'] + 1e-5)

# --- 时间跨度特征 ---
X_train['Promotion_Frequency'] = X_train['YearsSinceLastPromotion'] / (X_train['YearsAtCompany'] + 1)
X_test['Promotion_Frequency'] = X_test['YearsSinceLastPromotion'] / (X_test['YearsAtCompany'] + 1)

X_train['RoleChange_Frequency'] = X_train['YearsInCurrentRole'] / (X_train['TotalWorkingYears'] + 1)
X_test['RoleChange_Frequency'] = X_test['YearsInCurrentRole'] / (X_test['TotalWorkingYears'] + 1)

# --- 满意度组合 ---
satisfaction_cols = ['EnvironmentSatisfaction', 'JobSatisfaction', 'RelationshipSatisfaction']
X_train['Satisfaction_Mean'] = X_train[satisfaction_cols].mean(axis=1)
X_test['Satisfaction_Mean'] = X_test[satisfaction_cols].mean(axis=1)

X_train['Satisfaction_Std'] = X_train[satisfaction_cols].std(axis=1)
X_test['Satisfaction_Std'] = X_test[satisfaction_cols].std(axis=1)

# --- 对数变换减少偏度 ---
for col in ['MonthlyIncome', 'DistanceFromHome', 'NumCompaniesWorked']:
    if col in X_train.columns:
        X_train[col] = np.log1p(X_train[col])
        X_test[col] = np.log1p(X_test[col])

# 识别非数值型的列（即分类特征）
categorical_features = X_train.select_dtypes(include=['object']).columns
# 对分类特征进行独热编码
X_train_encoded = pd.get_dummies(X_train[categorical_features],dtype=int)
X_test_encoded = pd.get_dummies(X_test[categorical_features],dtype=int)
print(X_train_encoded)

# 接下来，我们需要合并独热编码的特征回原始的数值型特征中
# 首先，识别数值型的列
numerical_features = X_train.select_dtypes(exclude=['object']).columns

# 选择数值型特征
X_train_numerical = X_train[numerical_features]
X_test_numerical = X_test[numerical_features]

# 合并数值型特征和独热编码的特征
X_train = pd.concat([X_train_numerical, X_train_encoded], axis=1)
X_test = pd.concat([X_test_numerical, X_test_encoded], axis=1)

scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# 逻辑回归
lr = LogisticRegression(max_iter=50,random_state=22)
lr.fit(X_train, y_train)

os.makedirs("../model", exist_ok=True)
joblib.dump(lr, "../model/logistic_regression_model.pkl")
y_pred_lr = lr.predict_proba(X_test)[:,1]
print(f'Logistic Regression AUC: {roc_auc_score(y_test, y_pred_lr)}')
